import os
import glob
import subprocess
from Bio import SeqIO

"""
python collect_by_diamond.py
参数说明
脚本中的参数可以通过修改脚本中的变量来设置：

query_dir：查询文件所在的目录，默认为 "blast_seeds"。
fasta_dir：目标蛋白数据库文件所在的目录，默认为 "fasta_proteomes/example"。
result_dir：结果文件存放的目录，默认为 "mining_result"。
n：筛选出最优的 n 个结果，如果不指定则不筛选只排序，默认为 10。


fasta文件存放在fasta_proteomes/example/，用作query的蛋白文件存放在blast_seeds/，使用diamond, 将query文件与fasta_proteomes/example/中的所有fasta文件比对，将比对结果存放在比对结果文件夹mining_result中，比对结果文件名为diamond文件名+fasta文件名，并且取出所有匹配的蛋白全序列，保存为fasta文件。筛选其中最优的n个结果，如果不指定n，则不筛选只排序。
"""


def run_diamond(query_file, fasta_file, output_file):
    """Run Diamond BLASTP and save the results."""
    command = [
        "diamond", "blastp",
        "-q", query_file,
        "-d", fasta_file,
        "-o", output_file,
        "--outfmt", "6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore"
    ]
    subprocess.run(command, check=True)

def extract_matches(output_file, fasta_file, result_fasta):
    """Extract matching sequences from Diamond output and save to a new FASTA file."""
    with open(output_file, 'r') as f:
        matches = set(line.split()[1] for line in f)
    
    with open(fasta_file, 'r') as fasta, open(result_fasta, 'w') as out_fasta:
        for record in SeqIO.parse(fasta, 'fasta'):
            if record.id in matches:
                SeqIO.write(record, out_fasta, 'fasta')

def sort_and_filter_results(result_dir, n=None):
    """Sort and optionally filter the results based on bitscore."""
    for result_fasta in glob.glob(os.path.join(result_dir, '*_matches.fasta')):
        records = list(SeqIO.parse(result_fasta, 'fasta'))
        records.sort(key=lambda x: float(x.description.split()[-1]), reverse=True)
        
        if n is not None:
            records = records[:n]
        
        sorted_result_fasta = result_fasta.replace('_matches.fasta', '_sorted.fasta')
        with open(sorted_result_fasta, 'w') as out_fasta:
            SeqIO.write(records, out_fasta, 'fasta')
        
        os.remove(result_fasta)

def main(query_dir, fasta_dir, result_dir, n=None):
    # Create result directory if it doesn't exist
    os.makedirs(result_dir, exist_ok=True)
    
    # Get all query files
    query_files = glob.glob(os.path.join(query_dir, '*.fasta'))
    
    # Get all fasta files
    fasta_files = glob.glob(os.path.join(fasta_dir, '*.fasta'))
    
    for query_file in query_files:
        query_name = os.path.basename(query_file)
        
        for fasta_file in fasta_files:
            fasta_name = os.path.basename(fasta_file)
            
            # Define output file name
            output_file = os.path.join(result_dir, f'diamond_{query_name}_{fasta_name}.out')
            result_fasta = os.path.join(result_dir, f'{query_name}_{fasta_name}_matches.fasta')
            
            # Run Diamond
            run_diamond(query_file, fasta_file, output_file)
            
            # Extract matches
            extract_matches(output_file, fasta_file, result_fasta)
    
    # Sort and filter results
    sort_and_filter_results(result_dir, n)

if __name__ == "__main__":
    query_dir = "blast_seeds"
    fasta_dir = "fasta_proteomes/example"
    result_dir = "mining_result"
    n = 10  # Set to None if you don't want to filter
    
    main(query_dir, fasta_dir, result_dir, n)